import numpy as np
from mindspore import nn, context
from mindspore import Model
import mindspore.dataset as ds
from mindspore.profiler import Profiler
import mindspore.ops as ops
from mindspore import Tensor
import mindspore
from rfcn_head_torchvision import Rfcn_Head
from config import config

class Net(nn.Cell):
    def __init__(self):
        super(Net, self).__init__()
        self.rcnn_mask_squeeze, self.bbox_targets, self.rcnn_labels, self.con5 = self.build_param()
        self.head = Rfcn_Head(config, 2048, 2014, 7, 81, 2, 0.0625, 2)
        self.expend = ops.ExpandDims()
    def build_param(self):

        rcnn_mask_squeeze = np.full((1280,), False, dtype=bool)
        rcnn_mask_squeeze[:8] = True
        rcnn_labels = np.ones(1280, ) - 2
        rcnn_labels[0] = 4
        rcnn_labels[1] = 1
        rcnn_labels[2] = 2
        rcnn_labels[3] = 0
        rcnn_labels[4] = 77
        rcnn_labels[5] = 3
        rcnn_labels[6] = -1
        rcnn_labels[7] = 22
        rcnn_mask_squeeze = Tensor(rcnn_mask_squeeze, dtype=mindspore.bool_)
        bbox_targets = Tensor(np.random.rand(1280, 4), dtype=mindspore.float32)
        rcnn_labels = Tensor(rcnn_labels, dtype=mindspore.int32)
        # rois = Tensor(np.array([[0, 10, 10, 20, 20]] * 1280, dtype=mindspore.float32))
        conv5 = Tensor(np.random.rand(2, 2048, 80, 38), dtype=mindspore.float32)
        return rcnn_mask_squeeze,bbox_targets,rcnn_labels,conv5
    def construct(self, x):

        out = self.head(self.con5, x[0], self.bbox_targets, self.rcnn_labels, self.rcnn_mask_squeeze)
        out = self.expend(self.expend(out, 0), 0)
        return out

def generator():
    for i in range(2):
        yield (np.array([[[0, 10, 10, 20, 20]] * 1280]).astype(np.float32), np.ones(1).astype(np.int32))

def train(net):
    optimizer = nn.Momentum(net.trainable_params(), 1, 0.9)
    loss = nn.SoftmaxCrossEntropyWithLogits(sparse=True)
    data = ds.GeneratorDataset(generator, ["data", "label"])
    model = Model(net, loss, optimizer)
    model.train(1, data)

if __name__ == '__main__':
    # x = Tensor(np.ones([1280, 3969, 7, 7]), dtype=mindspore.float32)
    # print(ops.ReduceSum(keep_dims=False)(x,(2,3)).shape)
    # If the device_target is GPU, set the device_target to "GPU"
    # context.set_context(mode=context.GRAPH_MODE, device_target="Ascend", device_id = 3)
    context.set_context(mode=context.GRAPH_MODE, device_target="Ascend", device_id = 6)
    # Init Profiler
    # Note that the Profiler should be initialized after context.set_context and before model.train
    # If you are running in parallel mode on Ascend, the Profiler should be initialized before HCCL
    # initialized.
    profiler = Profiler(output_path="./rfcn_head_torchvison_train")

    # Train Model
    net = Net()
    net.set_train(True)
    #x = ops.Reshape()(ops.Tile()(Tensor(np.ones([1280, 7,7]),dtype=mindspore.float32),(15876,1 , 1)),(1280,15876,7,7))
    #x = Tensor(np.ones([1, 2]),dtype=mindspore.float32)
    # print(x)
    train(net)

    # Profiler end
    profiler.analyse()
